# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/23 22:16
# @Author : xyjiang1216
# @Email : jiangxuyan@nudt.edu.cn
# @File : construct_snowflake_topo_graph.py
# @Version: OpenPlanner2.0
import networkx as nx

from OpenPlanner.core.src.topo_construction_strategy.topo_interface import TopoStrategy
from OpenPlanner.util.topo_graph_utils import show_topology_graph


class SnowflakeTopoStrategy(TopoStrategy):
    # 重写方法
    def generate_topo_graph(self):
        # print('雪花形拓扑策略')
        # print(self.topo_generation_param.topo_type)

        sw_num = self.topo_param.sw_num
        es_num_per_sw_set = self.topo_param.es_num_per_sw_set
        if es_num_per_sw_set is None:
            print('配置文件中的es_num_per_sw_set不能为空！')
            exit(0)


        # 网络拓扑
        G = nx.DiGraph()

        # 第一步，添加交换机节点，以及交换机节点之间的链路
        sw_id = 0
        link_id = 0
        # 添加中心的交换机节点
        G.add_node(sw_id, node_id=sw_id, node_type='SW', es_set=[], current_avail_port=0)
        sw_id += 1

        former_node_num = 1
        last_layer_sw_id_set = []
        former_layer_sw_id_set = [0]
        cnt_layer_sw_id_set = []
        # 判断有几层
        # 如果只有一层，那么便可以开始添加终端节点
        if len(es_num_per_sw_set) == 1:
            last_layer_sw_id_set.append(0)
        # 如果大于一层，那么需要分层添加交换节点
        elif len(es_num_per_sw_set) >= 2:
            last_layer_sw_idx = 0
            for num_per_node in es_num_per_sw_set[:-1]:
                for former_sw_id in former_layer_sw_id_set:
                    for i in range(num_per_node):
                        G.add_node(sw_id, node_id=sw_id, node_type='SW', es_set=[], current_avail_port=0)
                        cnt_layer_sw_id_set.append(sw_id)
                        if last_layer_sw_idx == len(es_num_per_sw_set) - 2:
                            last_layer_sw_id_set.append(sw_id)

                        src_port = G.nodes[former_sw_id]['current_avail_port']
                        dst_port = G.nodes[sw_id]['current_avail_port']
                        G.add_edge(former_sw_id, sw_id, link_id=link_id, src_port=src_port, dst_port=dst_port)
                        link_id += 1
                        src_port = G.nodes[sw_id]['current_avail_port']
                        dst_port = G.nodes[former_sw_id]['current_avail_port']
                        G.add_edge(sw_id, former_sw_id, link_id=link_id, src_port=src_port, dst_port=dst_port)
                        link_id += 1
                        G.nodes[former_sw_id]['current_avail_port'] += 1
                        G.nodes[sw_id]['current_avail_port'] += 1

                        sw_id += 1
                former_layer_sw_id_set = cnt_layer_sw_id_set
                cnt_layer_sw_id_set = []
                last_layer_sw_idx += 1

        # 第二步，添加所有的终端节点
        # 以及终端节点与交换机节点之间的链路
        es_id = sw_id
        # 终端节点的数量
        es_node_num = former_node_num * es_num_per_sw_set[-1]
        # 终端节点应当连接到哪一个交换机
        for sw_id in last_layer_sw_id_set:
            for es in range(es_num_per_sw_set[-1]):
                G.add_node(es_id, node_id=es_id, node_type='ES', current_avail_port=0)

                src_port = G.nodes[sw_id]['current_avail_port']
                dst_port = G.nodes[es_id]['current_avail_port']
                G.add_edge(sw_id, es_id, link_id=link_id, src_port=src_port, dst_port=dst_port)
                link_id += 1
                src_port = G.nodes[es_id]['current_avail_port']
                dst_port = G.nodes[sw_id]['current_avail_port']
                G.add_edge(es_id, sw_id, link_id=link_id, src_port=src_port, dst_port=dst_port)
                link_id += 1
                G.nodes[es_id]['current_avail_port'] += 1
                G.nodes[sw_id]['current_avail_port'] += 1

                es_id += 1

        # 如果要显示拓扑
        show_topology_graph(G)

        # print(G)
        # 返回生成的拓扑图
        return G